package com.bw.test3;

import org.apache.flink.types.Row;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import com.alibaba.alink.pipeline.classification.KnnClassifier;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class KnnClassifierTest {
	@Test
	public void testKnnClassifier() throws Exception {
		BatchOperator.setParallelism(1);
		List <Row> df = Arrays.asList(
			Row.of(1, "0,0,0"),
			Row.of(1, "0.1,0.1,0.1"),
			Row.of(1, "0.2,0.2,0.2"),
			Row.of(0, "9,9,9"),
			Row.of(0, "9.1,9.1,9.1"),
			Row.of(0, "9.2,9.2,9.2")
		);
		BatchOperator <?> dataSource = new MemSourceBatchOp(df, "label int, vec string");

		KnnClassifier knn = new KnnClassifier().setVectorCol("vec")
			.setPredictionCol("pred")
				.setPredictionDetailCol("pred_detail")
			.setLabelCol("label")
			.setK(3);

		knn.fit(dataSource)
			.transform(dataSource)
			.print();
	}
}